import os
import shutil
import random

# 原始数据集路径
original_img_folder = r"D:\PyCharm\data\Thyroid\Thyroid_Dataset_all\tg3k\thyroid-image"
original_mask_folder = r"D:\PyCharm\data\Thyroid\Thyroid_Dataset_all\tg3k\thyroid-mask"

# 训练集和验证集保存路径
train_img_folder = r"D:\PyCharm\data\Thyroid\Thyroid_Dataset_all\tg3k\thyroid-train-image"
train_mask_folder = r"D:\PyCharm\data\Thyroid\Thyroid_Dataset_all\tg3k\thyroid-train-mask"
test_img_folder = r"D:\PyCharm\data\Thyroid\Thyroid_Dataset_all\tg3k\thyroid-test-image"
test_mask_folder = r"D:\PyCharm\data\Thyroid\Thyroid_Dataset_all\tg3k\thyroid-test-mask"

# 创建目标文件夹
os.makedirs(train_img_folder, exist_ok=True)
os.makedirs(train_mask_folder, exist_ok=True)
os.makedirs(test_img_folder, exist_ok=True)
os.makedirs(test_mask_folder, exist_ok=True)

# 获取所有图像文件列表
img_files = sorted(os.listdir(original_img_folder))

# 打乱文件列表
random.seed(42)  # 固定随机种子，确保每次运行结果一致
random.shuffle(img_files)

# 计算划分点
split_index = int(len(img_files) * 0.8)

# 划分为训练集和验证集
train_files = img_files[:split_index]
test_files = img_files[split_index:]

# 函数用于复制文件
def copy_files(file_list, src_folder, dest_folder):
    for file_name in file_list:
        src_path = os.path.join(src_folder, file_name)
        dest_path = os.path.join(dest_folder, file_name)
        shutil.copy(src_path, dest_path)

# 复制训练集图像和掩码
copy_files(train_files, original_img_folder, train_img_folder)
copy_files(train_files, original_mask_folder, train_mask_folder)

# 复制验证集图像和掩码
copy_files(test_files, original_img_folder, test_img_folder)
copy_files(test_files, original_mask_folder, test_mask_folder)

print("数据集已成功划分为训练集和验证集，并保存到指定文件夹。")